package com.tang.component.es;

import com.alibaba.fastjson2.JSON;
import lombok.extern.log4j.Log4j2;
import org.apache.commons.lang3.StringUtils;
import org.apache.http.HttpHost;
import org.elasticsearch.action.bulk.BulkRequest;
import org.elasticsearch.action.bulk.BulkResponse;
import org.elasticsearch.action.delete.DeleteRequest;
import org.elasticsearch.action.delete.DeleteResponse;
import org.elasticsearch.action.get.GetRequest;
import org.elasticsearch.action.get.GetResponse;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.action.index.IndexResponse;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.action.search.SearchType;
import org.elasticsearch.action.update.UpdateRequest;
import org.elasticsearch.action.update.UpdateResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestClientBuilder;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.common.unit.TimeValue;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.index.reindex.BulkByScrollResponse;
import org.elasticsearch.index.reindex.DeleteByQueryRequest;
import org.elasticsearch.search.Scroll;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.fetch.subphase.FetchSourceContext;
import org.elasticsearch.search.fetch.subphase.highlight.HighlightBuilder;
import org.elasticsearch.search.sort.ScoreSortBuilder;
import org.elasticsearch.search.sort.SortOrder;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.io.IOException;
import java.util.*;

@Log4j2
@Component
public class EsRestClient {

    //存放集群名称和集群连接客户端的映射关系
    private static final Map<String, RestHighLevelClient> clientMap = new HashMap<>();

    //配置信息
    @Resource
    private EsConfigurationProperties esConfigurationProperties;

    //公共请求项
    private static final RequestOptions COMMON_OPTIONS;

    static {
        RequestOptions.Builder builder = RequestOptions.DEFAULT.toBuilder();
        COMMON_OPTIONS = builder.build();
    }

    @PostConstruct
    protected void initialize() {
        List<EsClusterInfo> nodeInfoList = esConfigurationProperties.getNodeInfoList();
        try {
            for (EsClusterInfo esClusterInfo : nodeInfoList) {
                log.info("initialize.config.name:" + esClusterInfo.getName() + ",node:" + esClusterInfo.getNodes());
                RestHighLevelClient restHighLevelClient = initRestClient(esClusterInfo);
                String name = esClusterInfo.getName();
                clientMap.put(name, restHighLevelClient);
            }
        } catch (Exception e) {
            log.info("EsClient initError" + e);
        }
        log.info("initialize finish");
    }

    /**
     * 初始化es集群
     *
     * @param esClusterInfo es集群信息
     * @return org.elasticsearch.client.RestHighLevelClient
     * @author Tang
     * @date 2024/2/23 14:48
     **/
    private RestHighLevelClient initRestClient(EsClusterInfo esClusterInfo) {
        List<HttpHost> httpHostList = new ArrayList<>();
        String nodes = esClusterInfo.getNodes();
        String[] hosts = nodes.split(",");
        for (String host : hosts) {
            String[] ipPort = host.split(":");
            HttpHost httpHost = new HttpHost(ipPort[0], Integer.parseInt(ipPort[1]));
            httpHostList.add(httpHost);
        }
        HttpHost[] httpHosts = new HttpHost[httpHostList.size()];
        httpHostList.toArray(httpHosts);
        RestClientBuilder builder = RestClient.builder(httpHosts);
        return new RestHighLevelClient(builder);
    }

    /**
     * 根据集群名称获取集群连接客户端
     *
     * @param clusterName 集群名称
     * @return org.elasticsearch.client.RestHighLevelClient
     * @author Tang
     * @date 2024/2/23 14:52
     **/
    private static RestHighLevelClient getClient(String clusterName) {
        return clientMap.get(clusterName);
    }

    /**
     * 指定集群实例插入文档
     *
     * @param esIndexInfo  集群索引信息
     * @param esSourceData 需要插入的信息
     * @return boolean
     * @author Tang
     * @date 2024/2/23 14:51
     **/
    public static boolean insertDoc(EsIndexInfo esIndexInfo, EsSourceData esSourceData) {
        try {
            IndexRequest indexRequest = new IndexRequest(esIndexInfo.getIndexName());
            indexRequest.source(esSourceData.getData());
            indexRequest.id(esSourceData.getDocId());
            IndexResponse response = getClient(esIndexInfo.getNodeName()).index(indexRequest, COMMON_OPTIONS);
            if (log.isInfoEnabled()) {
                log.info("insertDoc.size:{}", JSON.toJSONString(response));
            }
            return true;
        } catch (Exception e) {
            log.error("insertDoc.exception:{}", e.getMessage(), e);
        }
        return false;
    }

    /**
     * 清空索引
     *
     * @param esIndexInfo 集群索引信息
     * @return boolean
     * @author Tang
     * @date 2024/2/23 15:07
     **/
    public static boolean delete(EsIndexInfo esIndexInfo) {
        try {
            DeleteByQueryRequest deleteByQueryRequest = new DeleteByQueryRequest(esIndexInfo.getIndexName());
            deleteByQueryRequest.setQuery(QueryBuilders.matchAllQuery());
            BulkByScrollResponse response = getClient(esIndexInfo.getNodeName()).deleteByQuery(deleteByQueryRequest, COMMON_OPTIONS);
            long deleted = response.getDeleted();
            if (log.isInfoEnabled()) {
                log.info("delete.size:{}", deleted);
            }
            return true;
        } catch (Exception e) {
            log.error("delete.exception:{}", e.getMessage(), e);
        }
        return false;
    }

    /**
     * 依据文档id和集群索引信息删除文档
     *
     * @param esIndexInfo 集群索引信息
     * @param docId       文档id
     * @return boolean
     * @author Tang
     * @date 2024/2/23 15:08
     **/
    public static boolean deleteDoc(EsIndexInfo esIndexInfo, String docId) {
        try {
            DeleteRequest deleteRequest = new DeleteRequest(esIndexInfo.getIndexName());
            deleteRequest.id(docId);
            DeleteResponse response = getClient(esIndexInfo.getNodeName()).delete(deleteRequest, COMMON_OPTIONS);
            if (log.isInfoEnabled()) {
                log.info("deleteDoc.size:{}", JSON.toJSONString(response));
            }
            return true;
        } catch (Exception e) {
            log.error("deleteDoc.exception:{}", e.getMessage(), e);
        }
        return false;
    }

    /**
     * 判断文档是否存在
     *
     * @param esIndexInfo 集群索引信息
     * @param docId       文档id
     * @return boolean
     * @author Tang
     * @date 2024/2/23 15:12
     **/
    public static boolean isExistDocById(EsIndexInfo esIndexInfo, String docId) {
        try {
            GetRequest getRequest = new GetRequest(esIndexInfo.getIndexName());
            getRequest.id(docId);
            return getClient(esIndexInfo.getNodeName()).exists(getRequest, COMMON_OPTIONS);
        } catch (Exception e) {
            log.error("isExistDocById.exception:{}", e.getMessage(), e);
        }
        return false;
    }


    /**
     * 根据文档id查询文档
     *
     * @param esIndexInfo 集群索引信息
     * @param docId       文档id
     * @return 集合
     * @author Tang
     * @date 2024/2/23 15:14
     **/
    public static Map<String, Object> getDocById(EsIndexInfo esIndexInfo, String docId) {
        try {
            GetRequest getRequest = new GetRequest(esIndexInfo.getIndexName());
            getRequest.id(docId);
            GetResponse response = getClient(esIndexInfo.getNodeName()).get(getRequest, COMMON_OPTIONS);
            if (log.isInfoEnabled()) {
                log.info(JSON.toJSONString(response));
            }
            return response.getSource();
        } catch (Exception e) {
            log.error("getDocById.exception:{}", e.getMessage(), e);
        }
        return null;
    }

    /**
     * 根据相关字段来搜索
     *
     * @param esIndexInfo 集群信息
     * @param docId       文档id
     * @param fields      需要查询的字段
     * @return java.util.Map<java.lang.String, java.lang.Object>
     * @author Tang
     * @date 2024/2/23 15:19
     **/
    public static Map<String, Object> getDocById(EsIndexInfo esIndexInfo, String docId, String[] fields) {
        try {
            GetRequest getRequest = new GetRequest(esIndexInfo.getIndexName());
            getRequest.id(docId);
            FetchSourceContext fetchSourceContext = new FetchSourceContext(true, fields, null);
            getRequest.fetchSourceContext(fetchSourceContext);
            GetResponse response = getClient(esIndexInfo.getNodeName()).get(getRequest, COMMON_OPTIONS);
            if (log.isInfoEnabled()) {
                log.info(JSON.toJSONString(response));
            }
            return response.getSource();
        } catch (Exception e) {
            log.error("getDocById.exception:{}", e.getMessage(), e);
        }
        return null;
    }

    /**
     * 全文检索
     *
     * @param esIndexInfo     集群索引信息
     * @param esSearchRequest es搜索条件构造
     * @return java.util.Map<java.lang.String, java.lang.Object>
     * @author Tang
     * @date 2024/2/23 15:28
     **/
    public static SearchResponse searchWithTermQuery(EsIndexInfo esIndexInfo, EsSearchRequest esSearchRequest) {
        try {
            BoolQueryBuilder bq = esSearchRequest.getBq();
            String[] fields = esSearchRequest.getFields();
            Integer from = esSearchRequest.getFrom();
            Boolean needScroll = esSearchRequest.getNeedScroll();
            Integer size = esSearchRequest.getSize();
            Long scrollCacheTime = esSearchRequest.getScrollCacheTime();
            String sortName = esSearchRequest.getSortName();
            SortOrder sortOrder = esSearchRequest.getSortOrder();
            HighlightBuilder highlighter = esSearchRequest.getHighlighter();
            //构建es搜索条件器
            SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
            searchSourceBuilder.query(bq);
            searchSourceBuilder.fetchSource(fields, null).from(from).size(size);
            //是否需要进行高亮处理
            if (Objects.nonNull(highlighter)) {
                searchSourceBuilder.highlighter(highlighter);
            }
            //是否指定了字段进行排序
            if (StringUtils.isNotBlank(sortName)) {
                searchSourceBuilder.sort(sortName);
            }
            //排序（相关度排序）默认从高到底
            searchSourceBuilder.sort(new ScoreSortBuilder().order(sortOrder));
            //构建请求对象
            SearchRequest searchRequest = new SearchRequest();
            searchRequest.searchType(SearchType.DEFAULT);
            searchRequest.indices(esIndexInfo.getIndexName());
            searchRequest.source(searchSourceBuilder);
            //是否开启快照
            if (needScroll) {
                Scroll scroll = new Scroll(TimeValue.timeValueMinutes(scrollCacheTime));
                searchRequest.scroll(scroll);
            }
            //发送请求
            SearchResponse response = getClient(esIndexInfo.getNodeName()).search(searchRequest, COMMON_OPTIONS);
            if (log.isInfoEnabled()) {
                log.info(JSON.toJSONString(response));
            }
            return response;
        } catch (Exception e) {
            log.error("searchWithTermQuery.exception:{}", e.getMessage(), e);
        }
        return null;
    }

    /**
     * 更新文档
     *
     * @param esIndexInfo  集群索引信息
     * @param esSourceData 文档数据
     * @return boolean
     * @author Tang
     * @date 2024/2/23 17:24
     **/
    public static boolean updateDoc(EsIndexInfo esIndexInfo, EsSourceData esSourceData) {
        try {
            UpdateRequest updateRequest = new UpdateRequest();
            updateRequest.index(esIndexInfo.getIndexName());
            updateRequest.id(esSourceData.getDocId());
            updateRequest.doc(esSourceData.getData());
            UpdateResponse response = getClient(esIndexInfo.getIndexName()).update(updateRequest, COMMON_OPTIONS);
            if (log.isInfoEnabled()) {
                log.info(JSON.toJSONString(response));
            }
            return true;
        } catch (Exception e) {
            log.error("updateDoc.exception:{}", e.getMessage(), e);
        }
        return false;
    }

    /**
     * 批量更新文档
     *
     * @param esIndexInfo      集群索引信息
     * @param esSourceDataList 文档列表
     * @return boolean
     * @author Tang
     * @date 2024/2/23 17:31
     **/
    public static boolean bathUpdateDoc(EsIndexInfo esIndexInfo, List<EsSourceData> esSourceDataList) {
        try {
            boolean flag = false;
            BulkRequest bulkRequest = new BulkRequest();
            for (EsSourceData esSourceData : esSourceDataList) {
                if (StringUtils.isBlank(esSourceData.getDocId())) {
                    continue;
                }
                flag = true;
                UpdateRequest updateRequest = new UpdateRequest();
                updateRequest.index(esIndexInfo.getIndexName());
                updateRequest.id(esSourceData.getDocId());
                updateRequest.doc(esSourceData.getData());
                bulkRequest.add(updateRequest);
            }
            if (flag) {
                BulkResponse response = getClient(esIndexInfo.getIndexName()).bulk(bulkRequest, COMMON_OPTIONS);
                if (log.isInfoEnabled()) {
                    log.info(JSON.toJSONString(response));
                }
                return !response.hasFailures();
            }
            return true;
        } catch (IOException e) {
            log.error("bathUpdateDoc.exception:{}", e.getMessage(), e);
        }
        return false;
    }
}
